# -*- coding: utf-8 -*-#
'''
# Name:         keras-fashionMNIST
# Description:  
# Author:       super
# Date:         2020/2/28
'''

import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'
import numpy as np
import pandas as pd
from contextlib import redirect_stdout
import matplotlib.pyplot as plt
from keras.datasets import fashion_mnist
from keras.utils import to_categorical
from keras.models import Sequential, load_model
from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from keras.optimizers import Adam

def load_data():
    (train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
    train_images = train_images.reshape(train_images.shape[0], 28, 28, 1)
    test_images = test_images.reshape(test_images.shape[0], 28, 28, 1)
    train_images = train_images.astype('float32') / 255
    test_images = test_images.astype('float32') / 255

    train_labels = to_categorical(train_labels, 10)
    test_labels = to_categorical(test_labels, 10)
    return (train_images, train_labels), (test_images, test_labels)

def build_model():
    model = Sequential()
    model.add(Conv2D(filters=16, kernel_size=2, padding='same', activation='relu', input_shape=(28, 28, 1)))
    model.add(MaxPooling2D(pool_size=2))

    model.add(Conv2D(filters=32, kernel_size=2, padding='same', activation='relu'))
    model.add(MaxPooling2D(pool_size=2))

    model.add(Conv2D(filters=64, kernel_size=2, padding='same', activation='relu'))
    model.add(MaxPooling2D(pool_size=2))
    model.add(Dropout(0.3))

    model.add(Flatten())
    model.add(Dense(500, activation='relu'))
    model.add(Dropout(0.4))
    model.add(Dense(10, activation='softmax'))

    adam = Adam(lr=1e-4)
    model.compile(optimizer=adam,
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])
    return model

def draw_train_history(history):
    plt.figure(1)

    # summarize history for accuracy
    plt.subplot(211)
    plt.plot(history.history['accuracy'])
    plt.plot(history.history['val_accuracy'])
    plt.title('model accuracy')
    plt.ylabel('accuracy')
    plt.xlabel('epoch')
    plt.legend(['train', 'test'], loc='upper left')

    # summarize history for loss
    plt.subplot(212)
    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('model loss')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.legend(['train', 'validation'], loc='upper left')
    plt.show()

def plot_predict(y_hat):
    fashion_mnist_labels = ["T-shirt/top",  # index 0
                            "Trouser",  # index 1
                            "Pullover",  # index 2
                            "Dress",  # index 3
                            "Coat",  # index 4
                            "Sandal",  # index 5
                            "Shirt",  # index 6
                            "Sneaker",  # index 7
                            "Bag",  # index 8
                            "Ankle boot"]  # index 9

    figure = plt.figure(figsize=(20, 8))
    for i, index in enumerate(np.random.choice(test_images.shape[0], size=15, replace=False)):
        ax = figure.add_subplot(3, 5, i + 1, xticks=[], yticks=[])
        # Display each image
        ax.imshow(np.squeeze(test_images[index]))
        predict_index = np.argmax(y_hat[index])
        true_index = np.argmax(test_labels[index])
        # Set the title for each image
        ax.set_title("{} ({})".format(fashion_mnist_labels[predict_index],
                                      fashion_mnist_labels[true_index]),
                     color=("green" if predict_index == true_index else "red"))
    plt.show()

if __name__ == '__main__':
    (train_images, train_labels), (test_images, test_labels) = load_data()
    model_path = "cnn_v1.h5"
    if os.path.exists(model_path):
        model = load_model(model_path)
    else:
        model = build_model()
        history = model.fit(train_images, train_labels, batch_size=64, epochs=11, validation_split=0.2)
        draw_train_history(history)
        model.save("cnn_v1.h5")

    with open('modelsummary.txt', 'w') as f:
        with redirect_stdout(f):
            model.summary()

    loss, accuracy = model.evaluate(test_images, test_labels)
    print("test loss: {}, test accuracy: {}".format(loss, accuracy))
    y_hat = model.predict(test_images)
    plot_predict(y_hat)